/*
 * tile_ints.h
 *
 *  Created on: Apr 18, 2013
 *      Author: drewlewis
 */

#ifndef TILE_PACKING_H_
#define TILE_PACKING_H_

#include "../Global.h"
#include <chemistry/molecule/molecule.h>
#include <chemistry/qc/basis/integral.h>
#include "IntegralEnginePool.h"
#include <boost/ref.hpp>
#include "integrals.hpp"
#include "range.hpp"
#include "tensor.hpp"
#include "tiled_array.h"
#include <Eigen/Dense>

namespace TA = TiledArray;

//Get all of the shells along one dimesion of a tile.
inline std::vector<int> get_shell_list(const atom_range &a, const Basis &basis){

  std::vector<int> a_shells;

  for(std::size_t i = a.first; i < a.second; ++i){
    std::size_t first_shell = basis->shell_on_center(i,0);
    std::size_t last_shell = basis->nshell_on_center(i) + first_shell;

    for(std::size_t j = first_shell; j < last_shell; ++j)
      a_shells.push_back(j);
  }

  return a_shells;
}

// Populates a tiles with integrals given a tile and integral object
template<typename RefEngine>
Tile2 get_ints(const ARRAY2 &a, const std::size_t i, const RefEngine &integral,
               const std::vector< atom_range > &atom_indices){

  TA::Range range = a.trange().make_tile_range(i);
  Tile2 tile(range); // output tile
  Basis basis = integral->basis();

  // Get list of shells for mpqc tensor object.
  std::vector<int> P = get_shell_list( atom_indices[0], basis);
  std::vector<int> Q = get_shell_list( atom_indices[1], basis);

  //Use a mpqc tensor to pack the tile
  const std::size_t dim[] = { range.size()[0], range.size()[1] };

  Tensor2 tile_map(tile.data(), dim);
  mpqc::integrals::evaluate(integral, P, Q, tile_map);

  return tile;
}

// Populates a tiles with integrals given a tile and integral object
template<typename RefEngine>
Tile3 get_ints(const ARRAY3 &a, const std::size_t i, const RefEngine &integral,
               const std::vector< atom_range > &atom_indices){

  TA::Range range = a.trange().make_tile_range(i);
  Tile3 tile(range); // output tile
  Basis basis_obs = integral->basis1();
  Basis basis_df = integral->basis3();

  // Get list of shells for mpqc tensor object.
  std::vector<int> P = get_shell_list( atom_indices[0], basis_obs);
  std::vector<int> Q = get_shell_list( atom_indices[1], basis_obs);
  std::vector<int> R = get_shell_list( atom_indices[2], basis_df);

  //Use a mpqc tensor to pack the tile
  const std::size_t dim[] = { range.size()[0], range.size()[1],
                              range.size()[2]};

  Tensor3 tile_map(tile.data(), dim);
  mpqc::integrals::evaluate(integral, P, Q, R, tile_map);

  return tile;
}

// Populates a tiles with integrals given a tile and integral object
template<typename RefEngine>
Tile4 get_ints(const ARRAY4 &a, const std::size_t i, const RefEngine &integral,
               const std::vector< atom_range > &atom_indices){

  TA::Range range = a.trange().make_tile_range(i);
  Tile4 tile(range); // output tile
  Basis basis = integral->basis();

  // Get list of shells for mpqc tensor object.
  std::vector<int> P = get_shell_list( atom_indices[0], basis);
  std::vector<int> Q = get_shell_list( atom_indices[1], basis);
  std::vector<int> R = get_shell_list( atom_indices[2], basis);
  std::vector<int> S = get_shell_list( atom_indices[3], basis);

  //Use a mpqc tensor to pack the tile
  const std::size_t dim[] = { range.size()[0], range.size()[1],
                              range.size()[2], range.size()[3]};

  Tensor4 tile_map(tile.data(), dim);
  mpqc::integrals::evaluate(integral, P, Q, R, S, tile_map);

  return tile;
}

// Task function to compute integrals for TA arrays
template<typename Eng_t, typename It, typename A>
void intgral_task( It first,  It last, const cluster_atoms &clusters, A &array,
                   boost::reference_wrapper< const sc::IntegralEnginePool<Eng_t> > pool,
                  const std::size_t block_size);

template <typename Eng_t, typename It, typename A>
void make_integral_task(It first, It last, const cluster_atoms &clusters,
                        A &array,
                        boost::reference_wrapper< const sc::IntegralEnginePool<Eng_t> > pool,
                        const std::size_t block_size) {

  array.get_world().taskq.add( &intgral_task<Eng_t, It, A>,  first,
                               last, clusters, array, pool, block_size);

}

// Task function to compute integrals for TA arrays
// Task function to compute integrals for TA arrays
template<typename Eng_t, typename It, typename A>
void intgral_task(It first, It last, const cluster_atoms &clusters, A &array,
                  boost::reference_wrapper<const sc::IntegralEnginePool<Eng_t> > pool,
                  const std::size_t block_size){

  while(last - first > block_size){
    It middle = first; // (last - fisrt)/2;
    std::advance( middle, std::distance(first, last)/2 );
    make_integral_task<Eng_t>(middle, last, clusters, array, pool, block_size);
    last = middle;
  }

  sc::Ref<Eng_t> integral = pool.get().intance();

  for(; first != last; ++first){

    //Get atoms on each tile
    const typename A::index index = array.range().idx(*first);
    std::vector< atom_range > atom_indices;
    for(std::size_t i = 0; i < index.size(); ++i)
      atom_indices.push_back( clusters[ index[i] ] );

    //Get the integrals for the tiles
    typename A::value_type tile = get_ints(array, *first,  integral,
                                               atom_indices);

    //Add tile to array
    array.set(*first, tile);
  }
}


// Fill up an array with tiles
template<typename Eng_t, class A>
void make_array_integrals(madness::World &world, A &array,
    Bases &bases, cluster_atoms &clusters, sc::IntegralEnginePool<Eng_t> &pool,
    std::size_t block_size = 1){

   typename A::pmap_interface::const_iterator  begin = array.get_pmap()->begin();
   typename A::pmap_interface::const_iterator end = array.get_pmap()->end();

   make_integral_task<Eng_t>(begin, end, clusters, array,
                      boost::cref(pool), block_size);
}

#endif /* TILE_PACKING_H_ */
